package com.mrx.commons.aop;

import com.mrx.commons.annotations.XRateLimit;
import com.mrx.commons.service.IRateLimiterService;
import com.mrx.commons.util.asserts.Asserts;
import org.apache.dubbo.config.annotation.DubboReference;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Optional;

/**
 * AOP 请求限流
 *
 * @author Mr.X
 * @since 2023-02-12 19:31
 */
@Aspect
@Component
public class AOPRateLimiter {

    private static final Logger logger = LoggerFactory.getLogger(AOPRateLimiter.class);

    @DubboReference
    private IRateLimiterService rateLimiterService;

    @Pointcut("@annotation(com.mrx.commons.annotations.XRateLimit)")
    public void rateLimiter() {
    }

    @Before("rateLimiter()")
    public void before(JoinPoint joinPoint) {
        Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
        XRateLimit rateLimiter = method.getAnnotation(XRateLimit.class);
        Optional.ofNullable(RequestContextHolder.getRequestAttributes())
                .map(it -> (ServletRequestAttributes) it)
                .map(ServletRequestAttributes::getRequest)
                .ifPresentOrElse(it -> handleRateLimit(rateLimiter, it, method.getName(), joinPoint.getArgs()), this::logError);
    }

    private void handleRateLimit(
            @NonNull XRateLimit rateLimiter, HttpServletRequest request,
            String methodName,
            Object... args
    ) {
        String rateLimitKey = ":" + methodName;
        if (rateLimiter.name() > -1) {
            rateLimitKey += ":" + args[rateLimiter.name()].toString();
        }
        long ttl = rateLimiterService.tryAcquire(request.getRemoteAddr() + rateLimitKey, rateLimiter.value());
        Asserts.RATE_LIMITED.checkTrue(ttl == IRateLimiterService.SUCCEED_LOCKED, String.format("请 %s 秒后再试", ttl));
    }

    private void logError() {
        logger.warn("request 为空, 限流失效!");
    }

}
